from typing_extensions import TypedDict
from typing import Annotated, Literal
from langgraph.graph.message import add_messages
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode
from langchain_community.chat_models import ChatTongyi
from langgraph.graph import StateGraph, END
# from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.memory import MemorySaver
from dotenv import load_dotenv
load_dotenv(verbose=True)

# LangGraph 持久化

class State(TypedDict):
    messages: Annotated[list, add_messages]


@tool
def search(query: str):
    """Call to surf the web."""
    # This is a placeholder for the actual implementation
    return ["The answer to your question lies within."]


tools = [search]
tool_node = ToolNode(tools)

# We will set streaming=True so that we can stream tokens
# See the streaming section for more information on this.
model = ChatTongyi()

bound_model = model.bind_tools(tools)

def should_continue(state: State) -> Literal["action", "__end__"]:
    """Return the next node to execute."""
    last_message = state["messages"][-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "__end__"
    # Otherwise if there is, we continue
    return "action"


def call_model(state: State):
    response = model.invoke(state["messages"])
    # We return a list, because this will get added to the existing list
    return {"messages": response}


# Define a new graph
workflow = StateGraph(State)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    # First, we define the start node. We use `agent`.
    # This means these are the edges taken after the `agent` node is called.
    "agent",
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
)

# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
# 常规边， action -> agent
workflow.add_edge("action", "agent")

# memory = SqliteSaver.from_conn_string(":memory:")
memory = MemorySaver()
print(memory)
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile(checkpointer=memory)

from langchain_core.messages import HumanMessage


def main():
    state = {}  # 将state改为字典,以thread_id作为键
    config = {"configurable": {"thread_id": "2"}}

    while True:
        user_input = input("User (输入quit或exit退出对话)：")
        if user_input.lower() in ["exit", "quit"]:
            break

        if '$' in user_input:
            try:
                question, thread_id = user_input.split('$')
                config["configurable"]["thread_id"] = thread_id.strip()
                question = question.strip()
                print(f"New config: {config}")
            except ValueError:
                print("输入格式错误,请使用 '问题$线程ID' 的格式。")
                continue
        else:
            question = user_input.strip()

        thread_id = config["configurable"]["thread_id"]
        if thread_id not in state:
            state[thread_id] = {"messages": []}  # 为新的thread_id初始化对话历史

        input_message = HumanMessage(content=question)
        state[thread_id]["messages"].append(
            input_message)  # 将新的用户输入追加到当前thread_id的对话历史中
        print(f"Current config: {config}")

        new_state = {
            "messages": state[thread_id]["messages"]
        }  # 使用当前thread_id的对话历史作为新状态
        for event in app.stream(new_state, config, stream_mode="values"):
            response_message = event["messages"][-1]
            response_message.pretty_print()
            state[thread_id]["messages"].append(
                response_message)  # 将新的AI回复保存回当前thread_id的对话历史中


if __name__ == "__main__":
    main()
